import random

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
from sklearn.manifold import TSNE

mpl.use("TkAgg")


def draw_sentences(sentence_embeddings: np.ndarray, labels: np.ndarray, filepath: str, show=False):
    X_tsne = TSNE(n_components=2, random_state=33).fit_transform(sentence_embeddings)

    font = {"color": "darkred",
            "size": 13,
            "family": "serif"}
    color_normal = mpl.colors.Normalize(vmin=0.5, vmax=17.5)
    num_list = [11, 1, 12, 10, 17, 8, 9, 13, 15, 2, 16, 7, 3, 4, 6, 14, 5]
    cmap = mpl.colors.ListedColormap([plt.cm.jet(color_normal(i)) for i in num_list])
    plt.figure(figsize=(13, 10))
    plt.scatter(X_tsne[:, 0], X_tsne[:, 1], c=labels, alpha=0.6, cmap=cmap, s=150)
    # plt.title("t-SNE", fontdict=font)
    cbar = plt.colorbar(ticks=range(1, 18))
    cbar.set_label(label='SDG ID', fontdict=font)
    plt.clim(0.5, 17.5)
    plt.tick_params(labelsize=20)
    plt.tight_layout()
    plt.savefig(filepath)
    if show:
        plt.show()


def draw_sentences_single(sentence_embeddings: np.ndarray, labels: np.ndarray, to_draw: np.array, filepath: str,
                          show=False):
    X_tsne = TSNE(n_components=2, random_state=33).fit_transform(sentence_embeddings)

    font = {"color": "darkred",
            "size": 13,
            "family": "serif"}
    color_normal = mpl.colors.Normalize(vmin=0.5, vmax=17.5)
    num_list = [11, 1, 12, 10, 17, 8, 9, 13, 15, 2, 16, 7, 3, 4, 6, 14, 5]
    cmap = mpl.colors.ListedColormap([plt.cm.jet(color_normal(i)) for i in num_list])
    fig = plt.figure(figsize=(13, 10))
    # ax = plt.subplot(1, 1, 1)
    label_names = [f"SGD{i}" for i in range(1, 18)]
    label_name_clsed = []
    label_clsed = []
    for i in to_draw:
        sentence_clsed = X_tsne[labels == i]
        clses = labels[labels == i]
        print(clses)
        label_name_clsed.append(label_names[i - 1])
        label_clsed.append(
            plt.scatter(sentence_clsed[:, 0], sentence_clsed[:, 1], c=[cmap(c) for c in clses], alpha=0.6, s=150))
    plt.tick_params(labelsize=20)
    # plt.title("t-SNE", fontdict=font)

    plt.legend(label_clsed, label_name_clsed, loc=1, bbox_to_anchor=(1.3, 1), prop={'size': 20})
    plt.tight_layout()
    plt.savefig(filepath)
    if show:
        plt.show()


if __name__ == '__main__':
    a = np.array([1, 2, 3, 4, 4, 5, 5])
    print(a[a == 4])
